from exception_hyd.embedding_exception import EmbeddingException
from numpy.typing import NDArray
from typing import List
import numpy as np
import pandas as pd
import requests
import os
import json


class Embedding:
    """
    创建一个embedding类
    """

    __private_api_token = ""
    __private_secret_token = ""

    def __init__(self, api_token: str = "", secret_token: str = "") -> None:
        self.__private_api_token = api_token
        self.__private_secret_token = secret_token

    def __get_access_token(self) -> str:
        if not self.__private_api_token:
            raise EmbeddingException("api token is empty")
        if not self.__private_secret_token:
            raise EmbeddingException("secret token is empty")
        if not (isinstance(self.__private_secret_token, str) and (self.__private_api_token, str)):
            raise EmbeddingException("api token and secret token must be 'str'")

        """
            使用 API Key Secret Key 获取access_token 替换下列示例中的应用API Key、应用Secret Key
        """

        url = f"https://aip.baidubce.com/oauth/2.0/token?" \
              f"grant_type=client_credentials&" \
              f"client_id={self.__private_api_token }&" \
              f"client_secret={ self.__private_secret_token }"

        payload = json.dumps("")
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        return str(response.json().get("access_token"))

    def get_embedding_data(self, text_list: List[str]) -> List[List[float]]:
        """
        该函数用于将文本列表embedding成向量矩阵

        :param text_list: 文本列表
        :return: 二维embedding矩阵
        """

        access_token = self.__get_access_token()
        url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1?access_token={access_token}"
        payload = json.dumps({
            "input": text_list
        })
        headers = {
            'Content-Type': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        response = json.loads(response.text)
        embedding_data = [i["embedding"] for i in response["data"]]
        return embedding_data

    def encode_str(self, query_list: List[str]) -> NDArray[np.float32]:
        """
        函数用于将查询的问题编码成向量

        :param query_list: 文本问题列表
        :return: 二维文本问题向量, 类型为np.ndarray, dtype=float32
        """

        access_token = self.__get_access_token()
        url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1?access_token={access_token}"
        payload = json.dumps({
            "input": query_list
        })
        headers = {
            'Content-Type': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        response = json.loads(response.text)
        embedding_data = [i["embedding"] for i in response["data"]]
        ndarray_float32_embedding_data = np.array(embedding_data, dtype="float32")

        # return ndarray[ndarray["float32"]]
        return ndarray_float32_embedding_data

    @staticmethod
    def save_embedding_data(embedding_data: List[List[float]], store_path_folder: str = "", file_name: str = "") -> str:
        """
        该函数用于将np.ndarray数据储存到硬盘上的指定位置

        :param embedding_data: 向量化后的文本向量
        :param store_path_folder: 储存文件夹，默认当ian文件夹
        :param file_name: 文件名，默认embedding_data.npy
        :return: 返回储存到硬盘的绝对文件地址
        """

        embedding_data = np.array(embedding_data, dtype="float32")
        if not store_path_folder:
            store_path_folder = os.getcwd()
        if not file_name:
            file_name = "embedding_data.npy"

        full_store_path = store_path_folder + "\\" + file_name
        np.save(full_store_path, embedding_data)
        return full_store_path

    @staticmethod
    def read_csv_col_tolist(csv_path: str, col_name: str) -> List[str]:
        """
        该函数用于读取csv文件的指定列

        :param csv_path: csv文件路径
        :param col_name: 列名
        :return: 返回一个文本列表
        """

        if not csv_path:
            raise EmbeddingException("the param ''csv_path' is empty")
        try:
            df = pd.read_csv(csv_path, encoding="gbk")

        except UnicodeDecodeError:
            df = df = pd.read_csv(csv_path, encoding="utf-8")

        return df[col_name].tolist()

    @staticmethod
    def read_excel_col_tolist(excel_path: str, col_name: "str") -> List[str]:
        """
         该函数用于读取xlsx文件的指定列

        :param excel_path: xlsx文件路径
        :param col_name: 列名
        :return: 返回一个文本列表
        """

        if not excel_path:
            raise EmbeddingException("the param ''csv_path' is empty")
        df = pd.read_excel(excel_path)

        return df[col_name].tolist()
